from flask import Flask, render_template, request, redirect, url_for, session, send_from_directory
import sqlite3
import re
import jieba
from werkzeug.security import generate_password_hash, check_password_hash

app = Flask(__name__)
app.secret_key = 'your-secret-key-here'
app.config['DATABASE'] = 'search_engine.db'

jieba.initialize()

# 中文停用词列表
STOP_WORDS = {
    '的', '是', '在', '了', '和', '有', '等', '我', '你', '他', '这', '那',
    '就', '也', '要', '不', '而', '与', '或', '为', '对', '但', '一个',
    '一些', '中', '地', '得', '了', '着', '过', '啊', '呀', '吗', '嗯'
}

# 同义词词典
SYNONYMS = {
    '电脑': ['计算机', 'PC'],
    '程序': ['程式', '软件'],
    '手机': ['移动电话', '智能手机'],
    '笔记本': ['笔记本电脑', '便携式电脑']
}

def get_db():
    conn = sqlite3.connect(app.config['DATABASE'])
    conn.execute("PRAGMA journal_mode = WAL")
    conn.execute("PRAGMA synchronous = NORMAL")
    conn.execute("PRAGMA cache_size = -1000000")
    conn.row_factory = sqlite3.Row
    return conn

def build_fts_query(raw_query):
    cleaned = re.sub(r'[^\w\u4e00-\u9fff\s-]', ' ', raw_query)
    has_chinese = any('\u4e00' <= c <= '\u9fff' for c in cleaned)
    
    # 分词并过滤停用词
    if has_chinese:
        terms = jieba.lcut_for_search(cleaned)
    else:
        terms = cleaned.split()
    terms = [term.strip() for term in terms if term.strip() and term not in STOP_WORDS]
    
    query_parts = []
    for term in terms:
        # 同义词扩展
        syn_terms = SYNONYMS.get(term, [])
        all_terms = [term] + syn_terms
        
        # 处理每个术语的格式
        term_group = []
        for t in all_terms:
            t = t.strip()
            if ' ' in t:
                term_group.append(f'"{t}"')
            elif t:
                term_group.append(t)
        
        if term_group:
            if len(term_group) > 1:
                query_parts.append(f'({" OR ".join(term_group)})')
            else:
                query_parts.append(term_group[0])
    
    return ' OR '.join(query_parts) if query_parts else ''

@app.route('/search')
def search():
    query = request.args.get('q', '').strip()
    page = int(request.args.get('page', 1))
    per_page = 10

    fts_query = build_fts_query(query)
    conn = get_db()
    
    try:
        where_clause = "pages_fts MATCH ?" if fts_query else "1=1"
        params = (fts_query, per_page, (page-1)*per_page) if fts_query else (per_page, (page-1)*per_page)

        # 调整BM25权重参数：假设标题权重为8.0，内容为1.0
        cursor = conn.execute(f"""
            SELECT rowid as id, title, 
            snippet(pages_fts, 2, '<b>', '</b>', '...', 64) as snippet,
            bm25(pages_fts, 8.0, 1.0) AS relevance
            FROM pages_fts 
            WHERE {where_clause}
            ORDER BY relevance DESC
            LIMIT ? OFFSET ?
        """, params)

        results = cursor.fetchall()
        total = conn.execute("SELECT COUNT(*) FROM pages_fts").fetchone()[0] if not fts_query else \
                conn.execute("SELECT COUNT(*) FROM pages_fts WHERE pages_fts MATCH ?", (fts_query,)).fetchone()[0]
    except sqlite3.OperationalError as e:
        print(f"Search error: {e}")
        results = []
        total = 0

    return render_template('results.html',
                         results=results,
                         query=query,
                         page=page,
                         total=total,
                         per_page=per_page)

@app.route('/admin/login', methods=['GET', 'POST'])
def admin_login():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        
        conn = get_db()
        user = conn.execute(
            'SELECT * FROM users WHERE username = ?', 
            (username,)
        ).fetchone()
        
        if user and check_password_hash(user['password'], password):
            session['admin_logged_in'] = True
            return redirect(url_for('admin_dashboard'))
    
    return render_template('admin/login.html')

@app.route('/admin/dashboard')
def admin_dashboard():
    if not session.get('admin_logged_in'):
        return redirect(url_for('admin_login'))
    
    conn = get_db()
    stats = conn.execute("""
        SELECT 
            (SELECT COUNT(*) FROM pages) AS page_count,
            (SELECT COUNT(*) FROM crawler_queue) AS queue_count
    """).fetchone()
    
    return render_template('admin/dashboard.html', stats=stats)

@app.route('/tencent8533094206726755089.txt')
def manifest():
    return send_from_directory('./', 'tencent8533094206726755089.txt')


@app.route('/')
def index():
    return render_template('search.html')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5200, debug=True)